package leetcode.biweekly.week82;

import java.util.Iterator;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.TreeMap;

//Solution3Test
public class Solution3 {
	// Solution3Test
	public long minSumSquareDiff(int[] nums1, int[] nums2, int k1, int k2) {
		TreeMap<Integer, Integer> tm = new TreeMap<>();
		for (int i = 0; i < nums1.length; i++) {
			int diff = Math.abs(nums1[i] - nums2[i]);
			int cnt = tm.getOrDefault(diff, 0);
			tm.put(diff, cnt + 1);
		}
		while (k1 > 0) {
			int cur = tm.lastKey();
			if (cur == 0) {
				return 0;
			}
			int cnt = tm.get(cur);
			tm.remove(cur);
			Integer next = tm.floorKey(cur - 1);
			next = next == null ? 0 : next;
			int diff = cur - next;
			if (k1 <= diff * cnt) {
				int d = k1 / cnt;
				int m = k1 % cnt;
				tm.put(cur - d, tm.getOrDefault(cur - d, 0) + cnt - m);
				tm.put(cur - d - 1, tm.getOrDefault(cur - d - 1, 0) + m);
				break;
			}
			int cnt2 = tm.getOrDefault(next, 0);
			tm.put(next, cnt2 + cnt);
			k1 -= diff * cnt;
		}
		while (k2 > 0) {
			int cur = tm.lastKey();
			if (cur == 0) {
				return 0;
			}
			int cnt = tm.get(cur);
			tm.remove(cur);
			Integer next = tm.floorKey(cur - 1);
			next = next == null ? 0 : next;
			int diff = cur - next;
			if (k2 <= diff * cnt) {
				int d = k2 / cnt;
				int m = k2 % cnt;
				tm.put(cur - d, tm.getOrDefault(cur - d, 0) + cnt - m);
				tm.put(cur - d - 1, tm.getOrDefault(cur - d - 1, 0) + m);
				break;
			}
			int cnt2 = tm.getOrDefault(next, 0);
			tm.put(next, cnt2 + cnt);
			k2 -= diff * cnt;
		}
		long ans = 0;
		for (Map.Entry<Integer, Integer> entry : tm.entrySet()) {
			ans += 1l * entry.getKey() * entry.getKey() * entry.getValue();
		}
		return ans;
	}

}
